package compute

import (
	"wgmat/gpucore"
)

func RegisterKernels(ctx gpucore.GPUContext) error {
	device := ctx.Device()

	if err := ctx.RegisterKernel("add", addWGSL, createAddLayout(device), [3]uint32{256, 1, 1}); err != nil {
		return err
	}

	if err := ctx.RegisterKernel("gemm", gemmWGSL, createGemmLayout(device), [3]uint32{16, 16, 1}); err != nil {
		return err
	}

	if err := ctx.RegisterKernel("lu_step", luKernelWGSL, createLUKernelLayout(device), [3]uint32{256, 1, 1}); err != nil {
		return err
	}
	layout := createLUSolveLayout(device)
	if err := ctx.RegisterKernel("forward_sub", forwardSubWGSL, layout, [3]uint32{1, 1, 1}); err != nil {
		return err
	}
	if err := ctx.RegisterKernel("backward_sub", backwardSubWGSL, layout, [3]uint32{1, 1, 1}); err != nil {
		return err
	}
	if err := ctx.RegisterKernel("find_pivot", findPivotWGSL, createFindPivotLayout(device), [3]uint32{256, 1, 1}); err != nil {
		return err
	}
	if err := ctx.RegisterKernel("swap_rows", swapRowsWGSL, createSwapRowsLayout(device), [3]uint32{256, 1, 1}); err != nil {
		return err
	}
	if err := ctx.RegisterKernel("lu_full", luFullWGSL, createLUFullLayout(device), [3]uint32{1, 1, 1}); err != nil {
		return err
	}
	if err := ctx.RegisterKernel("apply_pivot", applyPivotWGSL, createApplyPivotLayout(device), [3]uint32{256, 1, 1}); err != nil {
		return err
	}
	return nil
}
